package com.example.geoaimavlink_android.activity.utils;

import static androidx.core.math.MathUtils.clamp;

import android.graphics.Bitmap;
import android.graphics.ImageFormat;
import android.graphics.YuvImage;
import android.util.Log;

import androidx.annotation.NonNull;

import java.io.ByteArrayOutputStream;

public class YUVUtil {
    private static final String TAG = "YUVUtil";

    public static Bitmap yuv420ToBitmap(byte[] data, int width, int height) {
        YuvImage yuvImage = new YuvImage(data, ImageFormat.NV21, width, height, null);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        yuvImage.compressToJpeg(new android.graphics.Rect(0, 0, width, height), 100, outputStream);

        byte[] jpegData = outputStream.toByteArray();
        Bitmap bitmap = android.graphics.BitmapFactory.decodeByteArray(jpegData, 0, jpegData.length);

        try {
            outputStream.close();
        } catch (Exception e) {
            Log.e(TAG, "Error closing ByteArrayOutputStream: " + e.getMessage());
        }

        return bitmap;
    }

    public static byte[] convertI420toNV21(byte[] i420Data, int width, int height) {
        int ySize = width * height;
        int uvSize = ySize / 4;
        int nv21Size = ySize + 2 * uvSize;
        byte[] nv21Data = new byte[nv21Size];

        // 将Y平面数据复制到NV21数据
        System.arraycopy(i420Data, 0, nv21Data, 0, ySize);

        // 将U、V平面交错复制到NV21数据
        int uvIndex = ySize;
        for (int i = 0; i < uvSize; i++) {
            nv21Data[uvIndex++] = i420Data[ySize + i + uvSize];  // V
            nv21Data[uvIndex++] = i420Data[ySize + i];           // U
        }

        return nv21Data;
    }

    public static byte[] convertNV21toI420(byte[] nv21Data, int width, int height) {
        int frameSize = width * height;
        int chromaSize = frameSize / 4;
        int totalSize = frameSize + 2 * chromaSize;
        byte[] i420Data = new byte[totalSize];

        // Copy Y plane (luma)
        System.arraycopy(nv21Data, 0, i420Data, 0, frameSize);

        int offset = frameSize;
        int uIndex = frameSize;
        int vIndex = frameSize + chromaSize;

        // Deinterleave U and V planes (chroma)
        for (int i = 0; i < chromaSize; i++) {
            i420Data[uIndex + i] = nv21Data[offset++];  // U plane
            i420Data[vIndex + i] = nv21Data[offset++];  // V plane
        }

        return i420Data;
    }

    public static void conver_argb_to_i420(byte[] i420, int[] argb, int width, int height) {
        final int frameSize = width * height;

        int yIndex = 0;                   // Y start index
        int uIndex = frameSize;           // U statt index
        int vIndex = frameSize*5/4; // V start index: w*h*5/4

        int a, R, G, B, Y, U, V;
        int index = 0;
        for (int j = 0; j < height; j++) {
            for (int i = 0; i < width; i++) {
                a = (argb[index] & 0xff000000) >> 24; //  is not used obviously
                R = (argb[index] & 0xff0000) >> 16;
                G = (argb[index] & 0xff00) >> 8;
                B = (argb[index] & 0xff) >> 0;

                // well known RGB to YUV algorithm
                Y = ( (  66 * R + 129 * G +  25 * B + 128) >> 8) +  16;
                U = ( ( -38 * R -  74 * G + 112 * B + 128) >> 8) + 128;
                V = ( ( 112 * R -  94 * G -  18 * B + 128) >> 8) + 128;

                // I420(YUV420p) -> YYYYYYYY UU VV
                i420[yIndex++] = (byte) ((Y < 0) ? 0 : ((Y > 255) ? 255 : Y));
                if (j % 2 == 0 && i % 2 == 0) {
                    i420[uIndex++] = (byte)((U<0) ? 0 : ((U > 255) ? 255 : U));
                    i420[vIndex++] = (byte)((V<0) ? 0 : ((V > 255) ? 255 : V));
                }
                index ++;
            }
        }
    }

    public static byte[] bmp2Yuv(Bitmap bitmap) {
        int width = bitmap.getWidth();
        int height = bitmap.getHeight();

        int size = width * height;

        int pixels[] = new int[size];
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height);

        byte[] data = rgb2YCbCr420(pixels, width, height);
        return data;
    }

    public static byte[] rgb2YCbCr420(int[] pixels, int width, int height) {
        int len = width * height;
        // yuv格式数组大小，y亮度占len长度，u、v各占len/4长度
        byte[] yuv = new byte[len * 3 / 2];
        int y, u, v;
        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                // 屏蔽ARGB的透明度
                int rgb = pixels[i * width + j] & 0x00FFFFFF;
                // 像素的颜色顺序为bgr，移位运算
                int r = rgb & 0xFF;
                int g = (rgb >> 8) & 0xFF;
                int b = (rgb >> 16) & 0xFF;
                // 套用公式
                y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
                u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
                v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;
                // 调整
                y = y < 16 ? 16 : (y > 255 ? 255 : y);
                u = u < 0 ? 0 : (u > 255 ? 255 : u);
                v = v < 0 ? 0 : (v > 255 ? 255 : v);
                // 复制
                yuv[i * width + j] = (byte) y;
                yuv[len + (i >> 1) * width + (j & ~1) + 0] = (byte) u;
                yuv[len + +(i >> 1) * width + (j & ~1) + 1] = (byte) v;
            }
        }
        return yuv;
    }

    public static byte[] fetchNV21(@NonNull Bitmap bitmap) {
        int w = bitmap.getWidth();
        int h = bitmap.getHeight();
        int size = w * h;
        int[] pixels = new int[size];
        bitmap.getPixels(pixels, 0, w, 0, 0, w, h);


        byte[] nv21 = new byte[size * 3 / 2];

        // Make w and h are all even.
        w &= ~1;
        h &= ~1;


        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                int yIndex = i * w + j;

                int argb = pixels[yIndex];
                int a = (argb >> 24) & 0xff;  // unused
                int r = (argb >> 16) & 0xff;
                int g = (argb >> 8) & 0xff;
                int b = argb & 0xff;


                int y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
                y = clamp(y, 16, 255);
                nv21[yIndex] = (byte)y;

                if (i % 2 == 0 && j % 2 == 0) {
                    int u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
                    int v = ((112 * r - 94 * g -18 * b + 128) >> 8) + 128;


                    u = clamp(u, 0, 255);
                    v = clamp(v, 0, 255);


                    nv21[size + i / 2 * w + j] = (byte) v;
                    nv21[size + i / 2 * w + j + 1] = (byte) u;
                }
            }
        }
        return nv21;
    }
}
